#include "StdAfx.h"
#include "WinsockHook.h"
#include "Communication.h"

extern BOOL m_bIsConsoleCreated;
extern HANDLE m_console;

std::unordered_map<SOCKET, CSocketHelper*> CWinsockHook::mapSocketHelper;

INT WINAPI CWinsockHook::hook_connect(DWORD RetAddr, __pfnconnect pfnconnect, SOCKET s, const struct sockaddr FAR * name, int namelen)
{
#ifdef CREATE_CONSOLE
	if(m_bIsConsoleCreated == NULL)
	{
		AllocConsole();
		m_console = GetStdHandle(STD_OUTPUT_HANDLE);
		m_bIsConsoleCreated = TRUE;
	}
#endif

	Filter_connect(s, name, namelen);
	int ret = pfnconnect(s, name, namelen);

#ifdef CREATE_CONSOLE
// 	CHAR buffer[100] = {0};
// 	sprintf(buffer, "Connected. Result = %x.\n", ret);
// 	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

	// We assume all connections work
// 	if(ret == 0)
// 	{
		CSocketHelper *helper = new CSocketHelper(s);
		mapSocketHelper[s] = helper;
//	}
	return ret;
}

INT WINAPI CWinsockHook::hook_send(DWORD RetAddr, __pfnsend pfnsend, SOCKET s, const char FAR *buf, int len, int flags)
{
	CSocketHelper *helper = mapSocketHelper[s];
	if(helper == NULL)
	{
		CSocketHelper *helper = new CSocketHelper(s);
		mapSocketHelper[s] = helper;
	}
	if(len > 1)
	{
		Filter_send(s, buf, len, flags);
	}
	return pfnsend(s, buf, len, flags);
}

INT WINAPI CWinsockHook::hook_sendto(DWORD RetAddr, __pfnsendto pfnsendto, SOCKET s, const char* buf, int len, int flags, const struct sockaddr* to, int tolen)
{
	Filter_sendto(s, buf, len, flags, to, tolen);
	return pfnsendto(s, buf, len, flags, to, tolen);
}

INT WINAPI CWinsockHook::hook_recv(DWORD RetAddr, __pfnrecv pfnrecv, SOCKET s, char FAR *buf, int len, int flags)
{
#ifdef CREATE_CONSOLE
	if(m_bIsConsoleCreated == NULL)
	{
		AllocConsole();
		m_console = GetStdHandle(STD_OUTPUT_HANDLE);
		m_bIsConsoleCreated = TRUE;
	}
#endif

	INT ret;
	CSocketHelper *helper = mapSocketHelper[s];
	if(helper == NULL || !helper->IsCacheDataAvailable())
	{
		ret = pfnrecv(s, buf, len, flags);

#ifdef CREATE_CONSOLE
		CHAR buffer[100] = {0};
		sprintf(buffer, "socket = %d, received %d bytes data directly, WSALastError=0x%x.\n", s, ret, WSAGetLastError());
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

		if(WSAGetLastError() == 0 && ret > 1)
		{
			Filter_recv(s, buf, ret, flags);
		}
	}
	else
	{
// 		CHAR buff[12000];
// 		pfnrecv(s, buff, len, flags);

		// closesocket(s);
		ret = len;
		// Now we only retrieve cached data from main app

#ifdef CREATE_CONSOLE
// 		CHAR buffer[100] = {0};
// 		sprintf(buffer, "socket = %x, Fetching cached data.\n", s, len, ret);
// 		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

		ret = helper->GetCacheData((PUCHAR)buf, ret);

		if(ret == -1)
		{
			WSASetLastError(WSAEWOULDBLOCK);
		}
		else
		{
			WSASetLastError(0);
		}

#ifdef CREATE_CONSOLE
// 		sprintf(buffer, "socket = %d, receive %d bytes data from cache, WSALastError=0x%x\n", s, ret, WSAGetLastError());
// 		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
	}

// 	if(ret > 0)
// 	{
//  		CHAR strFilename[100] = {0};
//  		sprintf(strFilename, "D:\\sock-%d.dat", s);
//  		FILE *pf = fopen(strFilename, "ab");
// 		CHAR buffer[100] = {0};
// 		sprintf(buffer, "||%d-", ret);
// 		fwrite(buffer, strlen(buffer), 1, pf);
//  		fwrite(buf, ret, 1, pf);
//  		fclose(pf);
// 	}

	return ret;
}

INT WINAPI CWinsockHook::hook_recvfrom(DWORD RetAddr, __pfnrecvfrom pfnrecvfrom, SOCKET s, char* buf, int len, int flags, struct sockaddr* from, int* fromlen)
{
	INT ret;
	CSocketHelper *helper = mapSocketHelper[s];
	ret = pfnrecvfrom(s, buf, len, flags, from, fromlen);
	if(WSAGetLastError() == 0)
	{
		Filter_recvfrom(s, buf, ret, flags, from, fromlen);
	}
	return ret;
}

INT WINAPI CWinsockHook::hook_WSASend(DWORD RetAddr, __pfnWSASend pfnWSASend,
	SOCKET s, __in_ecount(dwBufferCount) LPWSABUF lpBuffers, 
	DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, 
	LPWSAOVERLAPPED lpOverlapped, 
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	INT ret = pfnWSASend(s, lpBuffers, dwBufferCount, 
		lpNumberOfBytesSent, dwFlags, lpOverlapped, lpCompletionRoutine);
	
	Filter_WSASend(s, lpBuffers, dwBufferCount, lpNumberOfBytesSent, dwFlags, lpOverlapped, lpCompletionRoutine);
	
	return ret;
}

INT WINAPI CWinsockHook::hook_WSARecv(
	DWORD RetAddr, 
	__pfnWSARecv pfnWSARecv,
	SOCKET s,
	LPWSABUF lpBuffers,
	DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesRecvd,
	LPDWORD lpFlags,
	LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine
	)
{
	CSocketHelper *helper = mapSocketHelper[s];
	
	Notice_WSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
	
	INT ret;

#ifdef CREATE_CONSOLE
	CHAR buffer[200] = {0};
	sprintf(buffer, "==========\nWSARECV === socket %x, hEvent = %x\n", 
		s, (lpOverlapped == NULL? 0 : lpOverlapped->hEvent));
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

	if(helper == NULL || !helper->IsCacheDataAvailable())
	{
#ifdef CREATE_CONSOLE
		sprintf(buffer, "Calling original WSARecv. \r\n");
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

		ret = pfnWSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
		if(lpOverlapped == NULL && lpCompletionRoutine == NULL)
		{
			// This is a blocked model!
			Filter_WSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
		}
		else if(lpOverlapped != NULL && lpOverlapped->hEvent != NULL)
		{
			if(helper != NULL)
			{
				helper->RegisterWSAEvent(lpOverlapped->hEvent, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, WHT_WSARECV);
				if(WSAGetLastError() == WSA_IO_PENDING)
				{
#ifdef CREATE_CONSOLE
					CHAR buffer[100] = {0};
					sprintf(buffer, "hook_WSARecv - IO_PENDING or %d. Register event %x.\n", WSAGetLastError(), lpOverlapped->hEvent);
					WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

					
				}
				else if(WSAGetLastError() == 0)
				{
#ifdef CREATE_CONSOLE
					CHAR buffer[100] = {0};
					sprintf(buffer, "hook_WSARecv - Got %d bytes data.\n", *lpNumberOfBytesRecvd);
					WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
					// We really get data
					Filter_WSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
// 					lpBuffers->buf[100] = 'Z';
// 					lpBuffers->buf[101] = 'Z';
				}
			}
		}
	}
	else
	{
#ifdef CREATE_CONSOLE
		sprintf(buffer, "lpOverlapped == %x, lpCompletionRoutine = %x.\r\n", lpOverlapped, lpCompletionRoutine);
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
		if(lpOverlapped == NULL && lpCompletionRoutine == NULL)
		{
			// This is in blocked mode
			// Read data directly
			ret = -1;
			while(ret == -1)
			{
				ret = helper->GetCacheData(lpBuffers, dwBufferCount);
				Sleep(1);
			}
			if(ret == 0)
			{
				// No more data
				*lpNumberOfBytesRecvd = 0;
				ret = 0;
			}
			else
			{
				// There returned %ret% bytes of data
				*lpNumberOfBytesRecvd = ret;
				ret = 0;
			}
		}
		else
		{
#ifdef CREATE_CONSOLE
			sprintf(buffer, "lpOverlapped->hEvent = %x.\r\n", lpOverlapped->hEvent);
			WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
			if(lpOverlapped != NULL && lpOverlapped->hEvent != NULL)
			{
// 				ret = pfnWSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
// 
#ifdef CREATE_CONSOLE
				CHAR buffer[100] = {0};
				sprintf(buffer, "hook_WSARecv - We have our own data.\n", ret, WSAGetLastError());
				WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

				// WSAEvent mode
				CWSABufferInfo* pBufInfo = helper->RegisterWSAEvent(lpOverlapped->hEvent, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, WHT_WSARECV);
				
				// Get our own data for acceleration!
				INT nGetCacheData = helper->GetCacheData(lpBuffers, dwBufferCount);
				if(nGetCacheData == -1)
				{
					// We should wait for more data
					// Start a new thread to set the signal
					helper->StartCacheStatusGuardingThread(lpOverlapped->hEvent);
					WSASetLastError(WSA_IO_PENDING);
					ret = SOCKET_ERROR;
					// *lpNumberOfBytesRecvd = 0;  - we should not update this according to MSDN
					pBufInfo->nOperationResult = SOCKET_ERROR;
				}
				else if(nGetCacheData == 0)
				{
					// No more data
					*lpNumberOfBytesRecvd = 0;
					ret = ERROR_SUCCESS;
					pBufInfo->nOperationResult = ERROR_SUCCESS;
					WSASetEvent(lpOverlapped->hEvent);
				}
				else
				{
// 					lpBuffers->buf[100] = 'B';
// 					lpBuffers->buf[101] = 'B';

					*lpNumberOfBytesRecvd = nGetCacheData;
					WSASetLastError(0);
					ret = ERROR_SUCCESS;
					pBufInfo->nOperationResult = ERROR_SUCCESS;
					pBufInfo->uNumberOfBytesRecvd = nGetCacheData;
					WSASetEvent(lpOverlapped->hEvent);
				}
			}
			else
			{
#ifdef CREATE_CONSOLE
				CHAR buffer[100] = {0};
				sprintf(buffer, "hook_WSARecv - Recv mode not supported.\n");
				WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
			}
		}
		
	}

#ifdef CREATE_CONSOLE
	sprintf(buffer, "hook_WSARecv - Exit. LastError=%d, recv bytes=%d, flags=%x.\n==========\n", WSAGetLastError(), *lpNumberOfBytesRecvd, *lpFlags);
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

	return ret;
}

INT WINAPI hook_WSARecvFrom(
	DWORD RetAddr, 
	__pfnWSARecvFrom pfnWSARecvFrom, 
	SOCKET s,                                               
	LPWSABUF lpBuffers,                                     
	DWORD dwBufferCount,                                    
	LPDWORD lpNumberOfBytesRecvd,                           
	LPDWORD lpFlags,                                        
struct sockaddr FAR *lpFrom,                            
	LPINT lpFromlen,                                        
	LPWSAOVERLAPPED lpOverlapped,                           
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine  
	)
{
	INT ret = pfnWSARecvFrom(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpFrom, lpFromlen, lpOverlapped, lpCompletionRoutine);
	Filter_WSARecvFrom(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpFrom, lpFromlen, lpOverlapped, lpCompletionRoutine);
	return ret;
}

INT WINAPI CWinsockHook::hook_WSAGetOverlappedResult(DWORD RetAddr, 
	__pfnWSAGetOverlappedResult pfnWSAGetOverlappedResult, 
	SOCKET s, LPWSAOVERLAPPED lpOverlapped, LPDWORD lpcbTransfer, BOOL fWait, LPDWORD lpdwFlags)
{
	CSocketHelper *helper = mapSocketHelper[s];
	INT ret;

#ifdef CREATE_CONSOLE
	CHAR buffer[200] = {0};
	sprintf(buffer, "OVERLAPPED === socket %x, hEvent = %x\n", 
		s, (lpOverlapped == NULL? 0 : lpOverlapped->hEvent));
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
	
	if(helper == NULL || !helper->IsCacheDataAvailable())
	{
		ret = pfnWSAGetOverlappedResult(s, lpOverlapped, lpcbTransfer, fWait, lpdwFlags);
		if(ret == TRUE)
		{
			CSocketHelper *helper = mapSocketHelper[s];
			if(helper != NULL && lpOverlapped->hEvent != NULL)
			{
				CWSABufferInfo* pBufferInfo = helper->GetWSABufferFromEvent(lpOverlapped->hEvent);
#ifdef CREATE_CONSOLE
				CHAR buffer[100] = {0};
				sprintf(buffer, "hook_WSAGetOverlappedResult - Get event %x, buffer ptr %x. Got bytes %d.\n", lpOverlapped->hEvent, pBufferInfo->lpBuffers, *lpcbTransfer);
				WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

				Notice_WSAGetOverlappedResult(s, lpOverlapped, lpcbTransfer, fWait, lpdwFlags, pBufferInfo);
// 				pBufferInfo->lpBuffers->buf[100] = 'X';
// 				pBufferInfo->lpBuffers->buf[101] = 'X';

				// helper->RemoveWSAEvent(lpOverlapped->hEvent);
				// delete pBufferInfo;
			}
			else
			{
				Notice_WSAGetOverlappedResult(s, lpOverlapped, lpcbTransfer, fWait, lpdwFlags, NULL);
			}
		}
	}
	else
	{
		// We have our own data for acceleration!
		if(lpOverlapped->hEvent != NULL)
		{
#ifdef CREATE_CONSOLE
			CHAR buffer[200] = {0};
			sprintf(buffer, "OVERLAPPED_OwnData === socket %x, hEvent = %x, helper = %x.\n", 
				s, (lpOverlapped == NULL? 0 : lpOverlapped->hEvent), helper);
			WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
			ret = pfnWSAGetOverlappedResult(s, lpOverlapped, lpcbTransfer, fWait, lpdwFlags);
			
			CWSABufferInfo* pBufferInfo = helper->GetWSABufferFromEvent(lpOverlapped->hEvent);

			
			if(pBufferInfo->nOperationResult == ERROR_SUCCESS)
			{
				// We have already post data to buffers
// 				pBufferInfo->lpBuffers->buf[100] = 'A';
// 				pBufferInfo->lpBuffers->buf[101] = 'A';
				*lpcbTransfer =pBufferInfo->uNumberOfBytesRecvd;
				ret = TRUE;

				// Remove registered event
				helper->RemoveWSAEvent(lpOverlapped->hEvent);
			}
			else
			{
				// We should get data from main app
				INT nBytesReceived = helper->GetCacheData(pBufferInfo->lpBuffers, pBufferInfo->dwBufferCount);
#ifdef CREATE_CONSOLE
				sprintf(buffer, "No more data. Get some from main app. \n");
				WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
				*lpcbTransfer = nBytesReceived;
				WSASetEvent(lpOverlapped->hEvent);
				helper->RemoveWSAEvent(lpOverlapped->hEvent);
				ret = TRUE;
			}
		}
		else
		{
			ret = FALSE;
		}
	}


	return ret;
}

BOOL WINAPI CWinsockHook::hook_CreateProcessAsUserW(DWORD retAddr,
	__pfnCreateProcessAsUserW pfnCreateProcessAsUserW,
	HANDLE hToken,
	LPCTSTR lpApplicationName,
	LPTSTR lpCommandLine,
	LPSECURITY_ATTRIBUTES lpProcessAttributes,
	LPSECURITY_ATTRIBUTES lpThreadAttributes,
	BOOL bInheritHandles,
	DWORD dwCreationFlags,
	LPVOID lpEnvironment,
	LPCTSTR lpCurrentDirectory,
	LPSTARTUPINFO lpStartupInfo,
	LPPROCESS_INFORMATION lpProcessInformation)
{
	BOOL ret = pfnCreateProcessAsUserW(hToken, lpApplicationName, lpCommandLine, 
		lpProcessAttributes, lpThreadAttributes, bInheritHandles, 
		dwCreationFlags, lpEnvironment,	lpCurrentDirectory, 
		lpStartupInfo, lpProcessInformation);

#ifdef CREATE_CONSOLE
	CHAR buffer[1024];
	sprintf(buffer, "[CreateProcessAsUser] PID = %d", lpProcessInformation->dwProcessId);
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
	// Hook the newly created process
	// TODO: Check whether it is a brower process
	CreateProcessInform(GetCurrentProcessId(), lpProcessInformation->dwProcessId);

	return ret;
}

CWinsockHook::CWinsockHook(void)
{
	pHookEnv_connect = InstallHook("ws2_32.dll", "connect", hook_connect);
	pHookEnv_send = InstallHook("ws2_32.dll", "send", hook_send);
	pHookEnv_sendto = InstallHook("ws2_32.dll", "sendto", hook_sendto);
	pHookEnv_recv = InstallHook("ws2_32.dll", "recv", hook_recv);
	pHookEnv_recvfrom = InstallHook("ws2_32.dll", "recvfrom", hook_recvfrom);
	pHookEnv_WSASend = InstallHook("ws2_32.dll", "WSASend", hook_WSASend);
	pHookEnv_WSARecv = InstallHook("ws2_32.dll", "WSARecv", hook_WSARecv);
	pHookEnv_WSARecvFrom = InstallHook("ws2_32.dll", "WSARecvFrom", hook_WSARecvFrom);
	pHookEnv_WSAGetOverlappedResult = InstallHook("ws2_32.dll", "WSAGetOverlappedResult", hook_WSAGetOverlappedResult);
	pHookEnv_CreateProcessAsUserW = InstallHook("Kernel32.dll", "CreateProcessAsUserW", hook_CreateProcessAsUserW);
}


CWinsockHook::~CWinsockHook(void)
{
	UninstallHook(pHookEnv_connect);
	UninstallHook(pHookEnv_send);
	UninstallHook(pHookEnv_sendto);
	UninstallHook(pHookEnv_recv);
	UninstallHook(pHookEnv_recvfrom);
	UninstallHook(pHookEnv_WSASend);
	UninstallHook(pHookEnv_WSARecv);
	UninstallHook(pHookEnv_WSARecvFrom);
	UninstallHook(pHookEnv_WSAGetOverlappedResult);
	UninstallHook(pHookEnv_CreateProcessAsUserW);
}
